# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import functools
import json
import logging

from coin.base.param_util import to_list
from coin.base.book.types import TradeSide
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.strategy.mm.feed import get_flow_book_builder_cls
from coin.feed.fastfeed import feed_pb2


class SingleProductGeneratorBase:
  def __init__(self):
    self._msgs = []

  def flush_fastfeed(self):
    msgs, self._msgs = self._msgs, []
    return msgs

  def _emit_proto_feed(self, msg):
    self._msgs.append(msg)


class SingleProductBookParamHolder(SingleProductGeneratorBase):
  def __init__(self, product, product_info=None):
    super().__init__()
    self._product = product
    self._symbol = str(self._product)

    product_info = product_info or get_holder_from_product(product).product_info
    self._price_precision = product_info.gen_tick_price().max_precision
    self._qty_precision = product_info.qty_precision

    self._price_multiplier = 10**self._price_precision
    self._qty_multiplier = 10**self._qty_precision

  def _refine_price_qty(self, price, qty):
    return round(price * self._price_multiplier), round(qty * self._qty_multiplier)


class SingleProductBookGenerator(SingleProductBookParamHolder):
  def __init__(self, product, product_info=None, feed_type=None, custom_feed_type=None):
    super().__init__(product, product_info=product_info)
    self._feed_type = feed_type or feed_pb2.FEED_TYPE_BOOK
    self._custom_feed_type = custom_feed_type

    self._ask_cleared = False
    self._ask_update = {}
    self._bid_cleared = False
    self._bid_update = {}

  def _reset_book_update(self):
    self._ask_cleared = False
    self._ask_update.clear()
    self._bid_cleared = False
    self._bid_update.clear()

  def on_ask_clear(self, book):
    assert self._product == book.product
    self._ask_cleared = True
    self._ask_update.clear()

  def on_ask_update(self, book, price, entry):
    assert self._product == book.product
    price, qty = self._refine_price_qty(price, entry.qty if entry else 0)
    self._ask_update[price] = qty

  def on_bid_clear(self, book):
    assert self._product == book.product
    self._bid_cleared = True
    self._bid_update.clear()

  def on_bid_update(self, book, price, entry):
    assert self._product == book.product
    price, qty = self._refine_price_qty(price, entry.qty if entry else 0)
    self._bid_update[price] = qty

  def on_book_built(self, book):
    assert self._product == book.product
    assert self._ask_cleared == self._bid_cleared

    ask_levels = []
    for price, qty in self._ask_update.items():
      try:
        ask_levels.append(feed_pb2.BookLevel(price=price, qty=qty))
      except ValueError:
        logging.exception('Error on %s ask book. Skipping the level. price=%d, price_precision=%d',
                          book.product,
                          price,
                          self._price_precision)

    bid_levels = []
    for price, qty in self._bid_update.items():
      try:
        bid_levels.append(feed_pb2.BookLevel(price=price, qty=qty))
      except ValueError:
        logging.exception('Error on %s bid book. Skipping the level. price=%d, price_precision=%d',
                          book.product,
                          price,
                          self._price_precision)

    msg = feed_pb2.Feed()
    msg.feed_type = self._feed_type
    if self._custom_feed_type:
      msg.custom_feed_type = self._custom_feed_type
    msg.symbol = self._symbol
    msg.timestamp = book.timestamp
    if book.exchange_timestamp:
      msg.origin_timestamp = book.exchange_timestamp
    if book.topic_id:
      msg.sort_key = book.topic_id

    if self._ask_cleared:
      msg.data_type = feed_pb2.DATA_TYPE_BOOK_SNAPSHOT
      msg.book_snapshot.ask.extend(ask_levels)
      msg.book_snapshot.bid.extend(bid_levels)
    else:
      msg.data_type = feed_pb2.DATA_TYPE_BOOK_UPDATE
      msg.book_update.ask.extend(ask_levels)
      msg.book_update.bid.extend(bid_levels)

    self._emit_proto_feed(msg)
    self._reset_book_update()


class SingleProductTradeGenerator(SingleProductBookParamHolder):
  def __init__(self, product, product_info=None, feed_type=None, custom_feed_type=None):
    super().__init__(product, product_info=product_info)
    self._feed_type = feed_type or feed_pb2.FEED_TYPE_BOOK
    self._custom_feed_type = custom_feed_type
    self._feed_continue_support = False

  def on_trade(self, trade):
    msg = feed_pb2.Feed()
    msg.feed_type = self._feed_type
    if self._custom_feed_type:
      msg.custom_feed_type = self._custom_feed_type
    msg.data_type = feed_pb2.DATA_TYPE_TRADE
    msg.symbol = self._symbol
    msg.timestamp = trade.timestamp

    price, qty = self._refine_price_qty(trade.price, trade.qty)
    msg.trade.price = price
    msg.trade.qty = qty
    if trade.side == TradeSide.TRADE_BUY_SIDE:
      msg.trade.side = feed_pb2.TRADE_SIDE_BUY
    elif (trade.side == TradeSide.TRADE_SELL_SIDE
          or trade.side == TradeSide.TRADE_SELL_SIDE_DEPRECATED):
      msg.trade.side = feed_pb2.TRADE_SIDE_SELL
    else:
      msg.trade.side = feed_pb2.TRADE_SIDE_UNKNOWN

    if trade.topic_id:
      msg.sort_key = trade.topic_id

    # TODO(jaewon): Implement from the parser level
    if trade.has_more_trade:
      # Infer whether the trade gen supports feed_continue.
      self._feed_continue_support = True
      msg.feed_continue = feed_pb2.FEED_CONTINUE_YES
    elif self._feed_continue_support:
      msg.feed_continue = feed_pb2.FEED_CONTINUE_NO

    self._emit_proto_feed(msg)


def _to_json_str(row):
  if isinstance(row, str):
    return row
  else:
    return json.dumps(row, separators=(',', ':'))  # Compact encoding


class SingleSymbolJsonObjectTableGenerator(SingleProductGeneratorBase):
  # We use the term `symbol` rather than `product` since it would handle non-product data.
  def __init__(self, symbol, feed_type, custom_feed_type=None):
    self._symbol = symbol.full_symbol
    self._feed_type = feed_type
    self._custom_feed_type = custom_feed_type
    self._rows = []

  def on_table_snapshot(self, timestamp):
    self.on_table_updated(timestamp, snapshot=True)

  def on_table_updated(self, timestamp, *, snapshot=False):
    msg = feed_pb2.Feed()
    msg.feed_type = self._feed_type
    if self._custom_feed_type:
      msg.custom_feed_type = self._custom_feed_type
    if snapshot:
      msg.data_type = feed_pb2.DATA_TYPE_JSON_OBJECT_TABLE_SNAPSHOT
    else:
      msg.data_type = feed_pb2.DATA_TYPE_JSON_OBJECT_TABLE_UPDATE
    msg.symbol = self._symbol
    msg.timestamp = timestamp
    msg.json_object_table = self._rows
    self._rows = []

    self._emit_proto_feed(msg)

  def on_row_inserted(self, key, row):
    row_msg = feed_pb2.JsonObjectTableRow()
    row_msg.action = feed_pb2.TABLE_ACTION_INSERT
    row_msg.key = key
    row_msg.json_object = _to_json_str(row)
    self._rows.append(row_msg)

  def on_row_updated(self, key, row):
    row_msg = feed_pb2.JsonObjectTableRow()
    row_msg.action = feed_pb2.TABLE_ACTION_UPDATE
    row_msg.key = key
    row_msg.json_object = _to_json_str(row)
    self._rows.append(row_msg)

  def on_row_deleted(self, key):
    row_msg = feed_pb2.JsonObjectTableRow()
    row_msg.action = feed_pb2.TABLE_ACTION_DELETE
    row_msg.key = key
    self._rows.append(row_msg)


def generate_feed_arb_param(timestamp, sub_req, product, product_info=None, sort_key=None):
  msg = feed_pb2.Feed()
  msg.feed_type = feed_pb2.FEED_TYPE_METADATA
  msg.data_type = feed_pb2.DATA_TYPE_FEED_ARB_PARAMETER
  msg.symbol = str(product)
  msg.timestamp = timestamp
  if sort_key:
    msg.sort_key = sort_key
  msg.feed_arb_param.symbol_group = get_flow_book_builder_cls(sub_req).find_feed_group(product)
  return msg


def generate_book_builder_param(timestamp, product, product_info=None, sort_key=None):
  msg = feed_pb2.Feed()
  msg.feed_type = feed_pb2.FEED_TYPE_METADATA
  msg.data_type = feed_pb2.DATA_TYPE_BOOK_BUILDER_PARAMETER
  msg.symbol = str(product)
  msg.timestamp = timestamp
  if sort_key:
    msg.sort_key = sort_key

  product_info = product_info or get_holder_from_product(product).product_info
  msg.book_builder_param.price_precision = product_info.gen_tick_price().max_precision
  msg.book_builder_param.qty_precision = product_info.qty_precision
  return msg


class ProtoFeedGeneratorBase:
  def __init__(self, feed_type, custom_feed_type=None):
    self._feed_type = feed_type
    self._custom_feed_type = custom_feed_type
    self._callbacks = []

  def publish_proto_feed(self, product, msgs):
    for callback in self._callbacks:
      callback(product, msgs)

  def subscribe_proto_feed(self, callback):
    self._callbacks.append(callback)


class BookGenerator(ProtoFeedGeneratorBase):
  def __init__(self, feed_type, custom_feed_type=None):
    super().__init__(feed_type, custom_feed_type)
    self._gen_map = {}

  def add_symbols(self, feed_handler, symbols):
    feed_handler.observe_book_builder(self)
    for symbol in to_list(symbols):
      feed_handler.subscribe_book(symbol, self.on_book)

  def _get_gen(self, product):
    try:
      return self._gen_map[product]
    except KeyError:
      gen = SingleProductBookGenerator(product,
                                       feed_type=self._feed_type,
                                       custom_feed_type=self._custom_feed_type)
      self._gen_map[product] = gen
      return gen

  def on_book(self, book):
    gen = self._get_gen(book.product)
    gen.on_book_built(book)
    self.publish_proto_feed(book.product, gen.flush_fastfeed())

  def on_ask_clear(self, book):
    self._get_gen(book.product).on_ask_clear(book)

  def on_ask_update(self, book, price, entry):
    self._get_gen(book.product).on_ask_update(book, price, entry)

  def on_bid_clear(self, book):
    self._get_gen(book.product).on_bid_clear(book)

  def on_bid_update(self, book, price, entry):
    self._get_gen(book.product).on_bid_update(book, price, entry)


class TradeGenerator(ProtoFeedGeneratorBase):
  def __init__(self, feed_type, custom_feed_type=None):
    super().__init__(feed_type, custom_feed_type)
    self._gen_map = {}

  def add_symbols(self, feed_handler, symbols):
    for symbol in to_list(symbols):
      feed_handler.subscribe_trade(symbol, functools.partial(self.on_trade, symbol))

  def _get_gen(self, product):
    try:
      return self._gen_map[product]
    except KeyError:
      gen = SingleProductTradeGenerator(product,
                                        feed_type=self._feed_type,
                                        custom_feed_type=self._custom_feed_type)
      self._gen_map[product] = gen
      return gen

  def on_trade(self, product, trade):
    gen = self._get_gen(product)
    gen.on_trade(trade)
    self.publish_proto_feed(product, gen.flush_fastfeed())


class JsonObjectTableGeneratorBase(ProtoFeedGeneratorBase):
  def __init__(self, feed_type, custom_feed_type=None):
    super().__init__(feed_type, custom_feed_type)
    self._gen_map = {}

  def _get_gen(self, symbol):
    try:
      return self._gen_map[symbol]
    except KeyError:
      gen = SingleSymbolJsonObjectTableGenerator(symbol,
                                                 feed_type=self._feed_type,
                                                 custom_feed_type=self._custom_feed_type)
      self._gen_map[symbol] = gen
      return gen

  def on_table_snapshot(self, symbol, timestamp):
    gen = self._get_gen(symbol)
    gen.on_table_snapshot(timestamp)

  def on_table_updated(self, symbol, timestamp):
    gen = self._get_gen(symbol)
    gen.on_table_updated(timestamp)

  def on_row_inserted(self, symbol, key, row):
    gen = self._get_gen(symbol)
    gen.on_row_inserted(key, row)

  def on_row_updated(self, symbol, key, row):
    gen = self._get_gen(symbol)
    gen.on_row_updated(key, row)

  def on_row_deleted(self, symbol, key):
    gen = self._get_gen(symbol)
    gen.on_row_deleted(key)


class BitmexInstrumentGenerator(ProtoFeedGeneratorBase):
  def __init__(self, feed_type=None, custom_feed_type=None):
    if feed_type is None and custom_feed_type is None:
      feed_type = feed_pb2.FEED_TYPE_CUSTOM
      custom_feed_type = 'bitmex_instrument'
    super().__init__(feed_type, custom_feed_type)

  def add_symbols(self, feed_handler, symbols):
    for symbol in to_list(symbols):
      feed_handler.subscribe_instrument(symbol, functools.partial(self.on_instrument, symbol))

  def on_instrument(self, product, instrument):
    msg = feed_pb2.Feed()
    msg.feed_type = self._feed_type
    if self._custom_feed_type:
      msg.custom_feed_type = self._custom_feed_type
    msg.data_type = feed_pb2.DATA_TYPE_JSON
    msg.symbol = str(product)
    msg.timestamp = instrument.timestamp
    if instrument.topic_id:
      msg.sort_key = instrument.topic_id
    msg.json = _to_json_str(instrument.inst)
    self.publish_proto_feed(product, [msg])


class OkexFuturesV3InstrumentGenerator(ProtoFeedGeneratorBase):
  def __init__(self, feed_type=None, custom_feed_type=None):
    if feed_type is None and custom_feed_type is None:
      feed_type = feed_pb2.FEED_TYPE_CUSTOM
      custom_feed_type = 'okex_instrument'
    super().__init__(feed_type, custom_feed_type)

  def add_symbols(self, feed_handler, symbols):
    for symbol in to_list(symbols):
      feed_handler.subscribe_instrument(symbol, functools.partial(self.on_instrument, symbol))

  def on_instrument(self, product, instrument):
    msg = feed_pb2.Feed()
    msg.feed_type = self._feed_type
    if self._custom_feed_type:
      msg.custom_feed_type = self._custom_feed_type
    msg.data_type = feed_pb2.DATA_TYPE_JSON
    msg.symbol = str(product)
    msg.timestamp = instrument.timestamp
    if instrument.topic_id:
      msg.sort_key = instrument.topic_id

    out = {}
    if instrument.funding_rate:
      out['funding_rate'] = instrument.funding_rate
    if instrument.mark_price:
      out['mark_price'] = instrument.mark_price
    if instrument.open_interest:
      out['open_interest'] = instrument.open_interest

    msg.json = _to_json_str(out)
    self.publish_proto_feed(product, [msg])
